# -*- coding: utf-8 -*-
# @Author : Gan
# @Time : 2024/6/13 9:42

import os
import cv2
import random
import numpy as np
from tqdm import tqdm
import math


def resize_image(img, scale_x, scale_y):
    new_width = int(img.shape[1] * scale_x)
    new_height = int(img.shape[0] * scale_y)
    resized_image = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
    return resized_image


def rotate_image(img, angle):
    (height, width) = img.shape[:2]

    # 计算旋转中心（图像中心）
    center = (width // 2, height // 2)
    # 计算旋转矩阵
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)

    # 计算旋转后图像的新边界尺寸，确保图像完整
    cos = np.abs(rotation_matrix[0, 0])
    sin = np.abs(rotation_matrix[0, 1])
    new_width = int((height * sin) + (width * cos))
    new_height = int((height * cos) + (width * sin))

    # 调整旋转矩阵，使旋转后的图像中心与原图中心对齐
    rotation_matrix[0, 2] += (new_width / 2) - center[0]
    rotation_matrix[1, 2] += (new_height / 2) - center[1]

    # 执行旋转并获取结果
    rotated_img = cv2.warpAffine(img, rotation_matrix, (new_width, new_height))

    return rotated_img


def shear_image(image, shear_factor_x, shear_factor_y):
    # 获取图像尺寸
    h, w = image.shape[:2]

    # 计算新的图像大小
    # 假设图像中心在变换后仍然位于中心
    # 对于x方向的shear，图像宽度会增加
    # 对于y方向的shear，图像高度会增加

    new_w = int(w + abs(shear_factor_x) * h)
    new_h = int(h + abs(shear_factor_y) * w)

    shear_matrix = np.float32([[1, shear_factor_x, 0],
                               [shear_factor_y, 1, 0],
                               [0, 0, 1]])

    # 取仿射矩阵的前两行用于warpAffine
    affine_matrix = shear_matrix[:2]
    if shear_factor_x < 0:
        affine_matrix[0, 2] = abs(shear_factor_x) * h
    if shear_factor_y < 0:
        affine_matrix[1, 2] = abs(shear_factor_y) * w

    # 应用仿射变换
    result = cv2.warpAffine(image, affine_matrix, (new_w, new_h))

    # 去除多余黑边
    gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    x, y, w, h = cv2.boundingRect(contours[0])
    result = result[y:y + h, x:x + w]
    return result


def add_gaussian_noise(image, mean=0, var=1):
    """
    向图像添加高斯噪声
    :param image: OpenCV图像对象（numpy数组）
    :param mean: 噪声的均值
    :param var: 噪声的方差
    :return: 带有噪声的OpenCV图像对象（numpy数组）
    """
    row, col, ch = image.shape
    gauss = np.random.normal(mean, var ** 0.5, (row, col, ch))
    gauss = gauss.reshape(row, col, ch)
    noisy = image + gauss

    # 确保像素值不会超出0到255的范围
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    return noisy


if __name__ == '__main__':
    crop_size_list = list(np.arange(250, 801, 50))
    crop_num_list = [6]
    rotate_angle_list = list(np.arange(-60, -19, 10)) + list(np.arange(20, 61, 10)) + [0]
    shear_factor = list(np.round(np.arange(-0.8, -0.29, 0.1), 1)) + list(np.round(np.arange(0.3, 0.81, 0.1), 1))
    scale_factor = list(np.round(np.arange(0.8, 1.51, 0.1), 1))
    crop_start = 0
    crop_end = 4000

    root_path = r'E:/datasets/satgeoloc_dataset\base_map'
    dst_path = r'E:/datasets/satgeoloc_dataset\query'
    os.makedirs(dst_path, exist_ok=True)
    img_list = os.listdir(root_path)
    # query_list = random.choices(img_list, k=309)
    for img_name in tqdm(img_list):
        img = cv2.imread(f'{root_path}/{img_name}', cv2.IMREAD_UNCHANGED)
        # random.seed(3407) 不固定随机种子，否则每张reference的crop_num，crop_size都是一样的
        crop_num = random.choice(crop_num_list)

        for i in range(crop_num):
            coord_h, coord_w = random.randint(crop_start, crop_end), random.randint(crop_start, crop_end)
            dh, dw = random.choice(crop_size_list), random.choice(crop_size_list)
            rotate_angle = random.choice(rotate_angle_list)
            shear_x, shear_y = random.choice(shear_factor), random.choice(shear_factor)
            scale_x, scale_y = random.choice(scale_factor), random.choice(scale_factor)

            result = img[coord_h:coord_h + dh, coord_w:coord_w + dw]
            result = resize_image(result, scale_x, scale_y)  # 缩放
            if shear_x * shear_y < 0:
                result = shear_image(result, shear_x, shear_y)  # shear
            else:
                result = rotate_image(result, rotate_angle)  # 旋转
            result = add_gaussian_noise(result)
            # cv2.imshow('Gaussian Noisy', result)  
            # cv2.waitKey(0)  
            # cv2.destroyAllWindows()
            cv2.imwrite(f'{dst_path}/{img_name[:-4]}_{coord_h}_{coord_w}_{coord_h + dh}_{coord_w + dw}.TIF', result)
